package top.zhangjianyong.blog.config.thread;

import static top.zhangjianyong.blog.constant.ThreadPoolConstant.*;

import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

/**
 * 线程池配置类
 *
 * @author zhangjianyong
 * @since 2024/8/7 上午9:28
 */
@Slf4j
@Configuration
public class ExecutorConfig {

  private static final int core = Runtime.getRuntime().availableProcessors();

  @Bean(name = ASYNC_SQL_EXECUTOR_BEAN)
  public Executor asyncSqlExecutor() {
    log.info("start asyncSqlExecutor");
    return createThreadPoolTaskExecutor(ASYNC_SQL_EXECUTOR_PREFIX);
  }

  @Bean(name = ASYNC_V_CODE_EXECUTOR_BEAN)
  public Executor asyncVCodeExecutor() {
    log.info("start asyncVCodeExecutor");
    return createThreadPoolTaskExecutor(ASYNC_V_CODE_EXECUTOR_PREFIX);
  }

  @Bean(name = ASYNC_VECTOR_STORE_EXECUTOR_BEAN)
  public Executor asyncVectorStoreExecutor() {
    log.info("start asyncVectorStoreExecutor");
    return createThreadPoolTaskExecutor(ASYNC_VECTOR_STORE_EXECUTOR_PREFIX);
  }

  private ThreadPoolTaskExecutor createThreadPoolTaskExecutor(String prefix) {

    ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
    // 配置核心线程数
    executor.setCorePoolSize(core);
    // 配置最大线程数
    executor.setMaxPoolSize(core * 2);
    // 配置队列大小
    executor.setQueueCapacity(50);
    // 配置线程池中的线程的名称前缀
    executor.setThreadNamePrefix(prefix);

    // rejection-policy：当pool已经达到max size的时候，如何处理新任务
    // CALLER_RUNS：不在新线程中执行任务，而是有调用者所在的线程来执行
    executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
    // 执行初始化
    executor.initialize();
    return executor;
  }
}
